import baostock as bs
import pandas as pd
from sqlalchemy.orm import Session

from models.bs_kline import BsKline
from mysql.database import get_db
from utils.stock import get_stock_list

def full_update_kline():
    stock_list=get_stock_list()
    stock_list = sorted(stock_list)
    #### 登陆系统 ####
    bs.login()
    for code in stock_list:
        if code > "002366.SZ":
            update_kline(code)
       
def update_kline(code:str):
    secid=code.split(".")[0]
    market=code.split(".")[1]
    bs_code=f"{market.lower()}.{secid}"
    rs = bs.query_history_k_data_plus(bs_code,
    "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",
    start_date='1990-12-19',
    frequency="d", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        row_data = rs.get_row_data()
        row_data[1] = code
        formatted_date = pd.to_datetime(row_data[0]).strftime("%Y%m%d")
        row_data[0]=int(formatted_date)
        data_list.append(row_data)
    df = pd.DataFrame(data_list, columns=rs.fields)
    db_gen = get_db()
    db: Session = next(db_gen)
    # existing_kline = db.query(BsKline).filter(BsKline.code == code).first()
    # if existing_kline:
    #    return 
 
    # 过滤空值
    df.dropna(inplace=True)

     # 数据清洗：将空字符串替换为 None，防止数据库插入错误
    numeric_columns = ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 
                      'turn', 'pctChg', 'peTTM', 'pbMRQ', 'psTTM', 'pcfNcfTTM']
    
    for col in numeric_columns:
        if col in df.columns:
            # 将空字符串替换为 None
            df[col] = df[col].replace('', None)
            # 同时移除仍然无法转换的行
            df = df[pd.to_numeric(df[col], errors='coerce').notna() | (df[col].isnull())]

    # 构造模型对象
    new_records = []
    # existing_dates = []

    # 查询已存在的 (code, date) 组合，避免重复插入
    # existing_date_set = set(
    #     db.query(BsKline.date).filter(BsKline.code == code).all()
    # )

    # existing_dates = {row[0] for row in existing_date_set}

    for _, row in df.iterrows():
        # date_int = int(row['date'])
        # if date_int in existing_dates:
        #     continue  # 跳过已存在的记录

        record = BsKline(
            **row
        )
        new_records.append(record)

    if not new_records:
        print(f"{code}无新K线数据需要插入")
        return 

    # 批量插入
    try:
        db.bulk_save_objects(new_records)
        db.commit()
        print(f"{code}查询到{len(data_list)}条记录，成功插入 {len(new_records)} 条K线数据")
       
    except Exception as e:
        db.rollback()
        print(f"{code}数据库插入失败: {str(e)}")
   